import gradio as gr, cv2, numpy as np, pickle
from PIL import Image

# 1. 载入模型
with open("best_knn.pickle", "rb") as f:
    knn = pickle.load(f)

def predict_digit(inp):
    # ===== ① 先看看到底有什么键 =====
    if not isinstance(inp, dict):
        return f"输入不是 dict，实际类型：{type(inp)}"
    print("[调试] inp.keys() =", list(inp.keys()))   # 终端会打印

    # ===== ② 万能取图：优先顺序 composite > image > layers =====
    if "composite" in inp:              # Gradio 4.x 最常见
        img = inp["composite"]
    elif "image" in inp:                # 旧版或某些模式
        img = inp["image"]
    elif "layers" in inp:               # 多层列表，取第一层
        img = inp["layers"][0] if len(inp["layers"]) else None
    else:
        return "找不到可用图字段"

    # ===== ③ 后续处理不变 =====
    gray = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY)
    gray = cv2.bitwise_not(gray)                # 白底黑字 → 黑底白字
    resized = cv2.resize(gray, (8, 8), interpolation=cv2.INTER_AREA)
    X = (resized / 16.0).reshape(1, -1)
    return str(knn.predict(X)[0])

# 4. 接口
with gr.Blocks(title="KNN 手写数字") as demo:
    gr.Markdown("### 最佳 KNN 手写数字识别")
    with gr.Row():
        pad = gr.Sketchpad(height=200, width=200, image_mode="RGB")
        lbl = gr.Label()
    btn = gr.Button("Submit")
    btn.click(predict_digit, inputs=pad, outputs=lbl)

demo.launch()